//
//  Segmenter.m
//  OneToOne
//
//  Created by ByteDance on 2023/3/10.
//

#import "Segmenter.h"
#import <GoogleMLKit/MLKit.h>
#import <CoreMedia/CoreMedia.h>
#import <CoreVideo/CoreVideo.h>
#import "UIUtilities.h"
static const CGFloat MLKSegmentationMaskAlpha = 1.0;
static int count = 1;
@implementation Segmenter {
    MLKSegmenter *segmenter;
}
- (instancetype)init {
    self = [super init];
    if (self) {
        MLKSelfieSegmenterOptions *options = [[MLKSelfieSegmenterOptions alloc] init];
        segmenter = [MLKSegmenter segmenterWithOptions:options];
    }
    return self;
}
- (void)processWithBuffer:(CMSampleBufferRef)sampleBuffer {
    CVImageBufferRef imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer);
    MLKVisionImage *visionImage = [[MLKVisionImage alloc] initWithBuffer:sampleBuffer];
    UIImageOrientation orientation = [UIUtilities
        imageOrientationFromDevicePosition:AVCaptureDevicePositionFront];
    visionImage.orientation = orientation;
    NSError *error;
    MLKSegmentationMask *mask = [segmenter resultsInImage:visionImage error:&error];
    CVPixelBufferRef gg = mask.buffer;
    if (imageBuffer) {
        if (mask != nil) {
                [UIUtilities applySegmentationMask:mask
                                     toImageBuffer:imageBuffer
                               withBackgroundColor:[UIColor.purpleColor colorWithAlphaComponent:MLKSegmentationMaskAlpha]
                                   foregroundColor:nil];

        } else {
          NSLog(@"Failed to segment image with error: %@", error.localizedDescription);
        }

        __weak __typeof(self) weakSelf = self;
        dispatch_sync(dispatch_get_main_queue(), ^{
          __strong typeof(weakSelf) strongSelf = weakSelf;
            if (count % 2 == 1) {
                [strongSelf.delegate receviceCVPixelBufferRef:imageBuffer];
            } else {
                [strongSelf.delegate receviceSendBuffer:imageBuffer];
            }
            count ++;
        });
    }
//    MLKVisionImage *image = [[MLKVisionImage alloc] initWithBuffer:sampleBuffer];
//    image.orientation =
//       [self imageOrientationFromDeviceOrientation:UIDevice.currentDevice.orientation
//                                    cameraPosition:AVCaptureDevicePositionFront];
//    __weak typeof(self) wself = self;
//    [segmenter processImage:image completion:^(MLKSegmentationMask * _Nullable mask,
//                                                  NSError * _Nullable error) {
//        __strong typeof(self) sself = wself;
//        CVImageBufferRef imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer);
//        if (mask != nil) {
//          UIColor *backgroundColor =
//              [UIColor.purpleColor colorWithAlphaComponent:MLKSegmentationMaskAlpha];
//          [UIUtilities applySegmentationMask:mask
//                               toImageBuffer:imageBuffer
//                         withBackgroundColor:backgroundColor
//                             foregroundColor:nil];
//        } else {
//          NSLog(@"Failed to segment image with error: %@", error.localizedDescription);
//        }
//        if ([sself.delegate respondsToSelector:@selector(receviceCVPixelBufferRef:)]) {
//            [sself.delegate receviceCVPixelBufferRef:imageBuffer];
//        }
//    }];
}
- (UIImageOrientation)
  imageOrientationFromDeviceOrientation:(UIDeviceOrientation)deviceOrientation
                         cameraPosition:(AVCaptureDevicePosition)cameraPosition {
  switch (deviceOrientation) {
    case UIDeviceOrientationPortrait:
      return cameraPosition == AVCaptureDevicePositionFront ? UIImageOrientationLeftMirrored
                                                            : UIImageOrientationRight;

    case UIDeviceOrientationLandscapeLeft:
      return cameraPosition == AVCaptureDevicePositionFront ? UIImageOrientationDownMirrored
                                                            : UIImageOrientationUp;
    case UIDeviceOrientationPortraitUpsideDown:
      return cameraPosition == AVCaptureDevicePositionFront ? UIImageOrientationRightMirrored
                                                            : UIImageOrientationLeft;
    case UIDeviceOrientationLandscapeRight:
      return cameraPosition == AVCaptureDevicePositionFront ? UIImageOrientationUpMirrored
                                                            : UIImageOrientationDown;
    case UIDeviceOrientationUnknown:
    case UIDeviceOrientationFaceUp:
    case UIDeviceOrientationFaceDown:
      return UIImageOrientationUp;
  }
}
@end
